import roboticstoolbox as rtb
import PyKDL as kdl
import numpy as np
from math import pi
import matplotlib.pyplot as plt
import sys
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')


class Motion_Planning:
    def __init__(self):
        # self.fingerU = rtb.models.URDF.IDMFingerURDF() # note
        self.finger = kdl.Chain()
        self.link_num = 5 # the number of link
        self.finger_dof = 4 # the number of joint
        
        jnts = [] # joint variable
        link = [] # link variable

        for i in range(self.link_num): 
            jnts.append(kdl.Joint(kdl.Joint.RotZ)) 
            
        frm1 = kdl.Frame(kdl.Rotation.RPY(0, 0, 0), kdl.Vector(0, 0.035, 0.05595))
        frm2 = kdl.Frame(kdl.Rotation.RPY(-pi / 2, pi / 2, pi), kdl.Vector(0, 0., 0))
        frm3 = kdl.Frame(kdl.Rotation.RPY(-pi / 2, pi / 6, pi), kdl.Vector(0, 0., 0))
        frm4 = kdl.Frame(kdl.Rotation.RPY(0, 0, 2.463 * pi / 180.), kdl.Vector(0.06832, 0.003, 0))
        frm5 = kdl.Frame(kdl.Rotation.RPY(0, 0, 0), kdl.Vector(0.033, -0.024, 0))

        frms = [frm1, frm2, frm3, frm4, frm5]

        for i in range(self.link_num):
            link.append(kdl.Segment(jnts[i], frms[i]))
            self.finger.addSegment(link[i])

        # the joint variable limints
        q_lim = np.array([[0, -pi / 4, np.deg2rad(-35), -pi / 2, -pi / 2],
                          [0, pi / 4, np.deg2rad(35), 0, np.deg2rad(6.)]])
        self.q_min = kdl.JntArray(self.link_num)
        self.q_max = kdl.JntArray(self.link_num)

        row, col = q_lim.shape
        for i in range(row):
            for j in range(col):
                if i == 0:
                    self.q_min[j] = q_lim[i][j]
                else:
                    self.q_max[j] = q_lim[i][j]
        
        self.fk = kdl.ChainFkSolverPos_recursive(self.finger)
    def fk_cal(self, qg):
        pos = kdl.Frame()
        qg_ = kdl.JntArray(self.link_num)
        qg_ = self.Real2KDL(qg, qg_)
        self.fk.JntToCart(qg_, pos)

        return qg_, pos
    
    def ik_cal(self, qg, path_num, rot_angle):
        qg_,pos= self.fk_cal(qg)
        Tran = pos.p
        ik_v = kdl.ChainIkSolverVel_pinv(self.finger)
        ik = kdl.ChainIkSolverPos_NR_JL(chain=self.finger, q_min=self.q_min, q_max=self.q_max, fksolver=self.fk, iksolver=ik_v, maxiter=100, eps=1e-6)

        q_out = kdl.JntArray(self.link_num)
        steps = self.cart_path(Tran, qg, path_num, rot_angle)
        sol = np.zeros((path_num, self.finger_dof))

        for i in range(path_num):
            for j in range(3):
                pos.p[j] = steps[i][j]
            ik.CartToJnt(q_init=qg_, p_in=pos, q_out=q_out)
            for j in range(self.finger_dof):
                sol[i][j] = q_out[j+1]

        qtraj = self.KDL2Real(sol)
        _, end_pos = self.fk_cal(qtraj[-1, :])

        return qtraj
    
    def jtraj_cal(self, start_q, end_q, nb_points):
        start_ = self.Real2Rtb(start_q)
        end_ = self.Real2Rtb(end_q)
        sol = rtb.jtraj(start_, end_, nb_points)
        jtraj = self.Rtb2Real(sol.q)

        return jtraj
    
    def jtraj_via_cal(self, start_q, via_q, via_qd, nb_points):
        start_q = self.Real2Rtb(start_q)
        row, col = via_q.shape
        s_points = np.sum(nb_points)
        sol_q = np.zeros((s_points, col))
        start_row = 0
        for i in range(row):
            via_q[i] = self.Real2Rtb(via_q[i])
            if i ==0:
                sol = rtb.jtraj(start_q, via_q[i], nb_points[i], qd1=via_qd[i])
                sol_q[start_row:nb_points[i], :] = sol.q
            else:
                sol = rtb.jtraj(via_q[i-1], via_q[i], nb_points[i], qd0=via_qd[i-1], qd1=via_qd[i])
                end_row = start_row + nb_points[i]
                sol_q[start_row:end_row, :] = sol.q
            start_row += nb_points[i]
        jtraj = self.Rtb2Real(sol_q)
        
        return jtraj
    
        
    def jac_cal(self, qg):
        qg_ = kdl.JntArray(self.link_num)
        qg_ = self.Real2KDL(qg, qg_)
        jac = kdl.ChainJntToJacSolver(self.finger)
        jac_matrix = kdl.Jacobian(self.link_num)
        jac.JntToJac(q_in=qg_, jac=jac_matrix)
        print(jac_matrix)
        return jac_matrix
            
    def Real2KDL(self, qg, qg_):
        for i in range(self.link_num):
            if i == 0:
                qg_[i] = 0
            elif i == 2:
                qg_[i] = qg[i-1]
            else:
                qg_[i] = -1 * qg[i-1]
        return qg_
    
    def KDL2Real(self, sol):
        qtraj = np.array(sol, copy=True)
        qtraj[:,0] = -1 * qtraj[:,0]
        qtraj[:,2] = -1 * qtraj[:,2]
        qtraj[:,3] = -1 * qtraj[:,3]
        qtraj = np.around(qtraj, 4)
        
        return qtraj
    
    def Real2Rtb(self, qg):
        qg_ = np.array(qg, copy=True)
        for i in range(self.finger_dof):
            if i == 1:
                qg_[i] = qg[i]
            else:
                qg_[i] = -1 * qg[i]
        
        return qg_
    
    def Rtb2Real(self, sol):
        qtraj = np.array(sol, copy=True)
        qtraj[:,0] = -1 * qtraj[:,0]
        qtraj[:,2] = -1 * qtraj[:,2]
        qtraj[:,3] = -1 * qtraj[:,3]
        qtraj = np.around(qtraj, 4)
        
        return qtraj
    
    def show(self, t, data):
        fig1 = plt.figure()
        ax1 = plt.subplot(4, 1, 1)
        plt.plot(t, data[:, 0], 'b', label='joint1')
        plt.legend()
        plt.grid()
        # plt.axis([0.0, 10, 0.0, 4.0])
        plt.subplot(4, 1, 2, sharex=ax1, sharey=ax1)
        plt.plot(t, data[:, 1], 'g', label='joint2')
        plt.legend()
        plt.grid()
        plt.subplot(4, 1, 3, sharex=ax1, sharey=ax1)
        plt.plot(t, data[:, 2], 'r', label='joint3')
        plt.legend()
        plt.grid()
        plt.subplot(4, 1, 4, sharex=ax1, sharey=ax1)
        plt.plot(t, data[:, 3], 'm', label='joint4')
        plt.legend()
        plt.grid()

        plt.show()
                
    

if __name__ == "__main__":
    demo = Motion_Planning()
    # qg = np.array(np.deg2rad([-4.,0.,28.,3.5]))
    # path_num = 10
    # rot_angle = 10.
    start_q = np.deg2rad([-4.,0.,28.,3.5])
    end_q = np.deg2rad([-10.,6.,25.,3.2])
    # via_q = np.deg2rad([[0.,0.,26., 1.],
    #                     [0.,0.,28.5, 2.5]])
    # via_qd = np.array([[0.05, 0.05, 0.05, 0.05],
    #                   [0., 0., 0., 0.]])
    nb_points = [5, 5] #第一段5个，第二段5个
    number_of_points = 10

    jtraj = demo.jtraj_cal(start_q, end_q, number_of_points)
    # jtraj = demo.jtraj_via_cal(start_q, via_q, via_qd, nb_points) #每段都要给个数
    s_points = np.sum(nb_points)
    t = np.linspace(0, s_points-1, s_points)

    print("Joint Trajectory (Rad):")
    for i in range(len(jtraj)):
        print(f"Time Step {i}: Joint Angles {np.rad2deg(jtraj[i])}")

    # demo.show(t, jtraj)
    

    
    




